import argparse
import os
from threading import Thread

import paramiko

cmd = "npu-smi info"

verbose = False
def ssh_execute_command(hostname, port, username, password, key_path, command):
    """
    通过SSH连接到远程主机并执行命令

    :param hostname: 远程主机IP
    :param port: SSH端口，默认22
    :param username: 登录用户名
    :param password: 登录密码
    :param command: 要执行的命令
    :return: 命令执行结果(stdout和stderr)
    """
    # 创建SSH客户端实例
    client = paramiko.SSHClient()

    try:
        # 自动添加主机密钥（生产环境中应考虑更安全的方式）
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        if key_path:
            private_key = paramiko.RSAKey.from_private_key_file(key_path)
            client.connect(hostname, port=port, username=username, pkey=private_key, timeout=5)
        elif password:
            client.connect(hostname, port=port, username=username, password=password, timeout=5)
        else:
            client.connect(hostname, port=port, username=username, timeout=5)

        # 执行命令
        stdin, stdout, stderr = client.exec_command(command)

        # 获取命令输出
        output = stdout.read().decode('utf-8').strip()
        error = stderr.read().decode('utf-8').strip()

        return {
            'output': output,
            'error': error,
            'status': stdout.channel.recv_exit_status()  # 获取命令返回状态码
        }

    except Exception as e:
        return {
            'output': '',
            'error': str(e),
            'status': -1
        }
    finally:
        # 关闭连接
        client.close()


def npu_status(host: str, username: str, password: str, pkey: str):
    ssh_resp = ssh_execute_command(host, 22, username, password, pkey, cmd)
    output = ssh_resp['output']
    output_line = output.split('\n')
    npu_status_list = []
    for item in output_line:
        if "No running processes found in NPU" in item:
            npu_status_list.append(item)
    if verbose:
        print(f"------------- {host} -------------")
        for npu_status in npu_status_list:
            print(f"unused npu: {npu_status}")
        print(f"-----------------------------------------")
    return npu_status_list


def check_npu_used(host: str, username: str, password: str, pkey: str, npu_count: int = 8):
    npu_status_list = npu_status(host, username, password, pkey)
    if len(npu_status_list) != npu_count:
        return False
    else:
        return True

def single_npu_result(host: str, username: str, password: str, pkey: str, npu_count: int, result_list: list):
    npu_status_list = npu_status(host, username, password, pkey)
    if len(npu_status_list) != npu_count:
        result_list.append({host: False})
    else:
        result_list.append({host: True})

def check_multi_npu_not_used(host_list: list, username: str, password: str, pkey: str, npu_count: int = 8):
    result_list = []
    thread_list = []
    for host in host_list:
        thread = Thread(target=single_npu_result, args=(host, username, password, pkey, npu_count, result_list))
        thread.start()
        thread_list.append(thread)
    for thread in thread_list:
        thread.join()

    multi_result = True
    idle_host_list = []
    busy_host_list = []
    for host in host_list:
        check_result = False
        for result_dict in result_list:
            if result_dict.get(host) is not None:
                check_result = result_dict.get(host)
                break
        print(f"host:{host} is not used:{check_result}")
        multi_result = multi_result and check_result
        if check_result:
            idle_host_list.append(host)
        else:
            busy_host_list.append(host)

    return multi_result, idle_host_list, busy_host_list

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # parser.add_argument('--host_list', required=False, type=str, default="192.168.0.94,192.168.0.120,192.168.0.73,192.168.0.118,192.168.0.58,192.168.0.212,192.168.0.4,192.168.0.31,192.168.0.147,192.168.0.211,192.168.0.182,192.168.0.223")
    parser.add_argument('--host_list', required=False, type=str, default="115.120.9.214,115.120.53.189,115.120.56.235")
    parser.add_argument('--username', required=False, type=str, default="root")
    parser.add_argument('--password', required=False, type=str, default="Campaign@2012")
    parser.add_argument('--pkey', required=False, type=str, default="")
    parser.add_argument('--npu_count', required=False, type=int, default=8)
    parser.add_argument('--verbose', action='store_true', help="verbose mode")
    args = parser.parse_args()
    verbose = args.verbose
    split_host_list = args.host_list.split(",")
    check_multi_npu_not_used(split_host_list, args.username, args.password, args.pkey)
